# -*- coding:utf-8 -*-

# @Time    : 2023/8/24 14:58
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : genarate_math_cot.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# -*- coding:utf-8 -*-

# @Time    : 2023/8/22 11:58
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : causal_dicover.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import aiohttp
import json
import asyncio
from agent.llm_agent import LLMAgent

async def call_gpt4(gpt_prompt):
    key_tuple = ("https://api.chatlingxi.com", "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")
    key = key_tuple[1]
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + key
    }
    payload = {
        "model": "gpt-4",
        "messages": gpt_prompt,
        "temperature": 0.9,
        "max_tokens": 2048,
        # "stream": True,
    }
    async with aiohttp.ClientSession() as session:
        response = await session.post(key_tuple[0] + "/v1/chat/completions",
                                      headers=headers, json=payload,
                                      timeout=aiohttp.ClientTimeout(total=100))
        content = await response.content.read()
    result = json.loads(content.decode())["choices"][0]['message']['content']
    return result

def read_data():
    # 读取jsonl文件
    # 文件内容的格式如下

    file_path = "../data_set/Math_data/math23k_train.json"
    text = open(file_path).read().replace('\n', '')
    lines = text.split("}")
    result = []
    for line in lines:
        if not line:
            continue
        try:
            json_data = json.loads(line + "}")
            new_json_data = {}
            new_json_data["问题"] = json_data["original_text"]
            new_json_data["方程式"] = json_data["equation"]
            new_json_data["答案"] = json_data["ans"]
            result.append(new_json_data)
        except:
            print(line)

    return result


class MathCot():
    def __init__(self):
        pass

    @classmethod
    async def generate(cls, context):
        context_str = "\n" + "问题:" + context["问题"] + "\n" + "方程式:" + str(context["方程式"]) + "\n" + "答案:" + str(context["答案"])
        gpt_prompt_str = """
        我希望你扮演一个“思维链的构建者”的角色。
你的任务是将给定的信息补充其推理的步骤用于通过COT（思维链）方式训练类型那些著名的AI系统（例如ChatGPT和GPT4）的推理能力。

但是完善后的思维链必须是合理的，且必须能被人类理解。
示例（参考推理逻辑，不用参考输出格式）
====
问题：吝啬鬼总喜欢把一条面包分成很多段吃，把一条面包分成5段用20秒，那么分成4段要用多少秒．
结果：该问题的重点是需要注意把一个面包分成5段是需要(5-1)即4次，而把面包分成4段则是(4-1)即3次。
因此，确定适用的数学公式：x=20/(5-1)*(4-1)。这个公式表示的是将一条面包分割的总时间（20秒）除以之前（5-1）次，然后乘以新的分割次数（4-1）次，得出结果：新的分割所需时间是15秒。所以，吝啬鬼把一条面包分成4段需要15秒。
=====
示例结束

信息如下：{context}

"##答案##"，”思维链”等术语不要出现在回复内容里面，也不要复述问题。在回复的思维链的内容里也不要包含"我们"或"我"等这种角色用语，只是针对问题参照方程式和答案，给出包含思维链的答案。
根据以上信息，完善思维链后的##答案##：
        """.format(context=context_str)
        llm_agent = LLMAgent(context["问题"], gpt_prompt_str)
        result = await llm_agent.achat_with_proxy_gpt4(save_path="cot")
        # 结果需要去掉"我们"、"思维链"等
        print(result)
        return result

if __name__ == '__main__':
    datas = read_data()
    datas = datas[5488:]
    for i in range(len(datas)):
        print(i)
        result = asyncio.run(MathCot.generate(datas[i]))
